from __future__ import division
from copy import copy
from pygame import draw, Surface, Rect

class Wave(object):
	def __init__(self, entry_node, timings, spawns):
		self.entry_node = entry_node
		self.start, self.end, self.delay = timings
		self.spawns = spawns
	
	def spawn(self, game, frame):
		if self.start <= frame and (frame < self.end or self.end == -1):
			if (frame - self.start) % self.delay == 0:
				for spawn in self.spawns:
					unit, deviation = spawn
					
					yield unit(game, self.entry_node, deviation)

class Graph:
	def __init__(self, waves):
		self.waves = waves
		
		self.start = min(wave.start for wave in waves)
		self.end = max(wave.end for wave in waves)
		self.duration = self.end - self.start
		
		if self.end != -1:
			self.wavey = {}
			for wave in self.waves:
				for frame in xrange(wave.start, wave.end, wave.delay):
					try:
						self.wavey[frame] += len(wave.spawns)
					except KeyError:
						self.wavey[frame] = len(wave.spawns)
						
			self.highest = max(self.wavey.values())
			
			self.graph = self.render(size=(1, 1))
		else:
			self.highest = 0

	def render(self, surface=None, size=None):
		if surface and not size:
			size = surface.get_size()
		elif size:
			surface = Surface(size).convert_alpha()
			surface.fill((0, 0, 0, 0))
		
		width, height = size
		
		for frame in self.wavey:
			frame_x = int(width * (frame / self.duration))
			
			bar_height = int(height * (self.wavey[frame] / self.highest))
			
			draw.line(surface, (255, 0, 0), (frame_x, 0), (frame_x, bar_height))
		
		return surface
	
	def draw(self, surface, rect, current_frame=0):
		if self.highest > 0:
			x, y = rect.topleft
			width, height = rect.width, rect.height
		
			if rect.size != self.graph.get_size():
				self.graph = self.render(size=rect.size)
		
			surface.blit(self.graph, (x, y))
		
			if self.start <= current_frame <= self.end:
				frame_x = int(x + width * (current_frame / self.duration))
			
				draw.line(surface, (255, 255, 255), (frame_x, y), (frame_x, y + height))
